package plgwsmsg

import (
	"dgo/goutils/utils"
	"fmt"
	"gitee.com/ymofen/gobase"
	"gitee.com/ymofen/gobase/quewriter"
	"gitee.com/ymofen/gobase/subpub"
	"golang.org/x/net/websocket"
	"io"
	"strings"
	"sync/atomic"
)

type WsMsgServ struct {
	online        int32
	closedflag    int32
	sessionSN     uint32
	sessionPrefix string

	wsMsgType byte

	ex *subpub.Subchannel

	handleRecvFunc func(conn *WsClientConn, buf []byte)

	OnError        func(conn *WsClientConn, err error)
	OnRecv         func(conn *WsClientConn, buf []byte)
	OnSent         func(conn *WsClientConn, buf []byte)
	OnConnected    func(conn *WsClientConn)
	OnDisconnected func(conn *WsClientConn, reason string)
}

type WsClientConn struct {
	closeflag  int32
	RemoteAddr string

	conn        *websocket.Conn
	subchannels map[string]struct{}
	SessionId   string
	TagPtr      interface{}
	disReason   string
	SubFunc     func(topic string)
	UnsubFunc   func(topic string) bool
	SendMsgFunc func(msg []byte)
}

func (this *WsClientConn) reqDis(s string) {
	if atomic.CompareAndSwapInt32(&this.closeflag, 0, 1) {
		this.disReason = s
		this.conn.Close()
	}
}

func ReceiveMessage(conn *websocket.Conn) (msgtype byte, msg []byte, err error) {
again:
	frame, err := conn.NewFrameReader()
	if err != nil {
		return 0, nil, err
	}
	frame, err = conn.HandleFrame(frame)
	if err != nil {
		return 0, nil, err
	}
	if frame == nil {
		goto again
	}
	maxPayloadBytes := conn.MaxPayloadBytes
	if maxPayloadBytes == 0 {
		maxPayloadBytes = websocket.DefaultMaxPayloadBytes
	}
	payloadType := frame.PayloadType()
	data, err := io.ReadAll(frame)
	if err != nil {
		return 0, nil, err
	}
	return payloadType, data, err
}

func SendWsMessage(conn *websocket.Conn, msgtype byte, msg []byte) error {
	if conn == nil {
		return io.ErrClosedPipe
	}
	writer, err := conn.NewFrameWriter(msgtype)
	if err != nil {
		return err
	}
	_, err = writer.Write(msg)
	return err
}

func NewWsMsgServ(sessionPrefix string) *WsMsgServ {
	return &WsMsgServ{sessionPrefix: sessionPrefix, ex: subpub.NewSubchannel()}
}

func (this *WsMsgServ) Online() int {
	return int(this.online)
}

func (this *WsMsgServ) StatusText() string {
	return this.ex.Status()
}

func (this *WsMsgServ) doErr(conn *WsClientConn, err error) {
	evt := this.OnError
	if evt != nil {
		evt(conn, err)
	}
}

func (this *WsMsgServ) doRecv(conn *WsClientConn, buf []byte) {
	evt := this.OnRecv
	if evt != nil {
		evt(conn, buf)
	}
}

func (this *WsMsgServ) doSent(conn *WsClientConn, buf []byte) {
	evt := this.OnSent
	if evt != nil {
		evt(conn, buf)
	}
}

func (this *WsMsgServ) doConnected(conn *WsClientConn) {
	evt := this.OnConnected
	if evt != nil {
		evt(conn)
	}
}

func (this *WsMsgServ) doDisconnected(conn *WsClientConn, reason string) {
	evt := this.OnDisconnected
	if evt != nil {
		evt(conn, reason)
	}
}

func (this *WsMsgServ) Close() error {
	atomic.StoreInt32(&this.closedflag, 1)
	this.ex.Pub("notify-close", 0, "服务关闭")
	return nil
}

func (this *WsMsgServ) HandleWsConn(conn *websocket.Conn) {
	if gobase.GoFunCatchException {
		defer gobase.DeferCatchPanic("HandleWsConn")
	}

	sn := atomic.AddUint32(&this.sessionSN, 1)
	sessionid := fmt.Sprintf("%s-%d", this.sessionPrefix, sn)

	wsconn := &WsClientConn{conn: conn, SessionId: sessionid, subchannels: make(map[string]struct{})}
	wsconn.RemoteAddr = conn.RemoteAddr().String()
	w0, err := conn.NewFrameWriter(this.wsMsgType)
	if err != nil {
		wsconn.reqDis(fmt.Sprintf("创建Writer异常:%s", err.Error()))
		return
	}
	defer w0.Close()
	quew := quewriter.NewQueueWriter(w0, 20480)

	quew.OnAfterWrite = func(p []byte, n int, err error) {
		if n > 0 {
			this.doSent(wsconn, p[:n])
		}
	}
	defer quew.Close()

	var sendfunc = func(msg []byte) {
		quew.Write(msg)
	}

	var subfunc = func(topic string) {
		if len(topic) == 0 {
			return
		}
		if _, ok := wsconn.subchannels[topic]; ok {
			return
		}

		this.ex.Sub(sessionid, topic, func(id, topic string, args ...interface{}) (ok bool) {
			if buf, ok := args[0].([]byte); ok {
				go sendfunc(buf)
			} else if str, ok := args[0].(string); ok {
				go sendfunc([]byte(str))
			}
			return true
		})
		wsconn.subchannels[topic] = struct{}{}
	}
	wsconn.SendMsgFunc = sendfunc
	wsconn.SubFunc = subfunc
	wsconn.UnsubFunc = func(topic string) bool {
		return this.ex.Unsub(sessionid, topic)
	}

	atomic.AddInt32(&this.online, 1)
	this.doConnected(wsconn)
	defer func() {
		atomic.AddInt32(&this.online, -1)
		this.doDisconnected(wsconn, wsconn.disReason)
	}()

	defer func() {
		for k, _ := range wsconn.subchannels {
			this.ex.Unsub(sessionid, k)
		}
		this.ex.Unsub(sessionid, "notify-close")
	}()

	if atomic.LoadInt32(&this.closedflag) == 1 {
		wsconn.reqDis("服务已经关闭")
		return
	}

	this.ex.Sub(sessionid, "notify-close", func(id, topic string, args ...interface{}) (ok bool) {
		wsconn.reqDis(args[0].(string))
		return true
	})

	str := conn.Request().URL.Query().Get("channels")
	if len(str) > 0 {
		strs := strings.Split(str, ",")
		for i := 0; i < len(strs); i++ {
			subfunc(utils.Trim(strs[i]))
		}
	}

	for {
		_, msg, err := ReceiveMessage(conn)
		if err != nil {
			wsconn.reqDis(err.Error())
			break
		}
		this.handleRecvFunc(wsconn, msg)
		this.doRecv(wsconn, msg)
	}
}
